use std::fs::File;
use std::io::{self, Read, Write};
use std::path::Path;
use md5::{Context};
use walkdir::WalkDir;
use sha2::{Sha256, Digest};  // 使用 Sha256 而非 Md5

fn calculate_md5<P: AsRef<Path>>(path: P) -> io::Result<String> {
    let mut file = File::open(path)?;
    let mut context = Context::new();
    let mut buffer = [0u8; 4096];

    loop {
        let bytes_read = file.read(&mut buffer)?;
        if bytes_read == 0 {
            break;
        }
        context.consume(&buffer[..bytes_read]);
    }

    Ok(format!("{:x}", context.compute()))
}

fn generate_md5_report(dir: &str) -> io::Result<()> {
    let output_file_path = "md5_report.txt";
    let mut output_file = File::create(output_file_path)?;

    for entry in WalkDir::new(dir) {
        let entry = entry?;
        let path = entry.path();

        if path.is_file() {
            let file_name = path.to_string_lossy();
            let md5 = calculate_md5(&path)?;
            writeln!(output_file, "{}: {}", file_name, md5)?;
            println!("Processed: {}", file_name);
        }
    }

    println!("MD5 report saved to {}", output_file_path);
    Ok(())
}


fn calculate_sha256<P: AsRef<Path>>(path: P) -> io::Result<String> {
    let mut file = File::open(path)?;
    let mut hasher = Sha256::new();  // 使用 Sha256
    let mut buffer = [0u8; 4096];

    loop {
        let bytes_read = file.read(&mut buffer)?;
        if bytes_read == 0 {
            break;
        }
        hasher.update(&buffer[..bytes_read]);  // 更新哈希值
    }

    Ok(format!("{:x}", hasher.finalize()))  // 计算最终的 SHA-256 哈希
}

fn generate_sha256_report(dir: &str) -> io::Result<()> {
    let output_file_path = "sha256_report.txt";
    let mut output_file = File::create(output_file_path)?;

    for entry in WalkDir::new(dir) {
        let entry = entry?;
        let path = entry.path();

        if path.is_file() {
            let file_name = path.to_string_lossy();
            let sha256 = calculate_sha256(&path)?;
            writeln!(output_file, "{}: {}", file_name, sha256)?;
            println!("Processed: {}", file_name);
        }
    }

    println!("SHA-256 report saved to {}", output_file_path);
    Ok(())
}


fn main() -> io::Result<()> {
    // let current_dir = ".";
    // generate_md5_report(current_dir);

    let current_dir = ".";
    generate_sha256_report(current_dir)
}